import math
import random
import pygame
import numpy as np
import rospy
import rvo.rvo_math as rvo_math
from rvo.vector import Vector2
from rvo.simulator_corridor import Simulator
from rvo import env
from matplotlib import rcParams
from draw.plt2d import plt_visulazation

RVO_RENDER = False


def transform_position(position, width, height, bounds):
    # Transform the position to fit within the window bounds.
    x_min, x_max, y_min, y_max = bounds

    x = (position.x - x_min) / (x_max - x_min) * width
    y = (position.y - y_min) / (y_max - y_min) * height

    return x, height - y  # Flip y-axis for correct display


class Blocks:
    def __init__(self):
        self.goals_ = []  # Vector2
        self.obstacles_ = []  # Vector2
        self.simulator_ = Simulator()
        self.agent_trajectories = {}
        self.previous_positions = {}

    def setup_scenario(self):
        # Specify the global time step of the simulation.
        self.simulator_.set_time_step(0.2)

        # Specify the default parameters for agents that are subsequently added.
        # neighborDist, maxNeighbors, timeHorizon, timeHorizonObst, radius, maxSpeed, velocity
        self.simulator_.set_agent_defaults(30.0, 10, 10.0, 10.0, 5.0, 5.0, Vector2(0.0, 0.0))

        # Add agents, specifying their start position, and store their goals on the opposite side of the environment.
        self.set_agents()
        self.agent_global_path_planning()

        # Add (polygonal) obstacles, specifying their vertices in counterclockwise order.
        self.set_obstacles()

        # Process the obstacles so that they are accounted for in the simulation.
        self.simulator_.process_obstacles()

        self.agent_trajectories = {i: [] for i in range(self.simulator_.num_agents)}
        self.previous_positions = {i: None for i in range(self.simulator_.num_agents)}

    def set_agents(self):
        starts = [[785.0, 12.0], [785.0, 27.0], [785.0, 42.0], [785.0, 57.0], [785.0, 72.0], [785.0, 87.0],
                  [785.0, 102.0], [785.0, 117.0], [785.0, 132.0], [785.0, 147.0], [785.0, 162.0], [785.0, 177.0]]
        goals = [[15, 12.0], [15, 27.0], [47, 42.0], [47, 57.0], [47, 72.0], [15, 87.0], [15, 102.0], [15, 117.0],
                 [15, 132.0], [70, 147.0], [70, 162.0], [70, 177.0]]
        for i in range(len(starts)):
            self.simulator_.add_agent(Vector2(starts[i][0], starts[i][1]), Vector2(goals[i][0], goals[i][1]))
            self.goals_.append(Vector2(goals[i][0], goals[i][1]))

            self.simulator_.add_agent(Vector2(goals[i][0], goals[i][1]), Vector2(starts[i][0], starts[i][1]))
            self.goals_.append(Vector2(starts[i][0], starts[i][1]))

    def set_obstacles(self):
        from rvo.config import envs_type
        environ = env.Env(envs_type[3])
        poly_obs = environ.poly_obs
        for i in range(len(poly_obs)):
            obstacle_i = []
            for j in range(len(poly_obs[i].vertices_pos)):
                obstacle_i.append(Vector2(poly_obs[i].vertices_pos[j][0], poly_obs[i].vertices_pos[j][1]))
            self.simulator_.add_obstacle(obstacle_i, polygon_id=poly_obs[i].id)
            self.obstacles_.append(obstacle_i)
        # print(len(self.simulator_.obstacles_))

    def agent_global_path_planning(self):
        environ = env.Env("large_corridor")
        obstacles = environ.poly_obs
        x_range, y_range = (-300, environ.xDim + 300), (0, environ.yDim)
        paths_smooth = []
        for i in range(self.simulator_.num_agents):
            agent = self.simulator_.agents_[i]
            radius = agent.radius_
            inflation = min(radius / 5, 1.2)
            all_path = [(self.goals_[i].x, self.goals_[i].y, 0)]
            agent.planner.init_rob_info((0, 0), (0, 0), radius, inflation=inflation)      # inflation用于膨胀障碍物，
            agent.planner.init_env(obstacles, x_range, y_range)
            for j in range(1):
                # 起止位置必须是tuple型
                s_start = (self.goals_[i].x, self.goals_[i].y)        # path是逆序，此处起止位置互换正好不用再次逆序
                s_goal = (agent.position_.x, agent.position_.y)
                agent.planner.init_rob_info(s_start, s_goal, radius, inflation=inflation)
                path = agent.planner.global_search()
                path[-1] = (path[-1][0], path[-1][1], 0)
                path.pop(0)
                all_path += path
            agent.path_ = all_path
            path_smooth = agent.path_
            paths_smooth.append(path_smooth)
        obstacles = obstacles
        # plt_visulazation(paths_smooth, self.env_.agents_, [], obstacles, task_area, exit_area)

    def update_visualization(self, screen, width, height, bounds):
        screen.fill((255, 255, 255))
        colors = rcParams['axes.prop_cycle'].by_key()['color']  # 获取默认颜色循环
        for i in range(self.simulator_.num_agents):
            position = self.simulator_.agents_[i].position_
            self.agent_trajectories[i].append(position)  # 记录轨迹
            transformed_radius = (width * self.simulator_.default_agent_.radius_) / (bounds[1] - bounds[0])
            transformed_position = transform_position(position, width, height, bounds)
            color = [0, 0, 0]
            color[i % 3] = 255  # Set one of the color channels to max (255) based on the agent index
            pygame.draw.circle(screen, color, transformed_position, transformed_radius)
            # 绘制智能体编号
            font = pygame.font.Font(None, 15)
            text = font.render(str(i), True, color)
            screen.blit(text, (transformed_position[0]-transformed_radius, transformed_position[1]-2*transformed_radius))
            # 绘制智能体轨迹
            # if len(self.agent_trajectories[i]) > 1:
            #     for j in range(1, len(self.agent_trajectories[i])):
            #         start_pos = transform_position(self.agent_trajectories[i][j - 1], width, height, bounds)
            #         end_pos = transform_position(self.agent_trajectories[i][j], width, height, bounds)
            #         pygame.draw.line(screen, color, start_pos, end_pos, 1)
        for idx, obstacle in enumerate(self.obstacles_):
            color = colors[idx % len(colors)]
            points = [transform_position(vec, width, height, bounds) for vec in obstacle]
            pygame.draw.polygon(screen, color, points, 0)  # Fill the polygon with black color

    def set_preferred_velocities(self):
        """
        Set the preferred velocity to be a vector of unit magnitude (speed) in the direction of the goal.
        """
        collision_num = 0
        collision_agent = []
        for i in range(self.simulator_.num_agents):
            goal_vector = self.goals_[i] - self.simulator_.agents_[i].position_
            distSq = rvo_math.abs_sq(goal_vector)
            # print('agent' + str(self.simulator_.agents_[i].id_) + ' 终点距离: ' + str(math.sqrt(distSq)))
            if self.simulator_.agents_[i].collision_:
                collision_num += 1
                collision_agent.append(i)
        # print("collided num: ", collision_num, "collided ids: ", collision_agent)

    def reached_goal(self):
        """
        Check if all agents have reached their goals.
        """
        for i in range(self.simulator_.num_agents):
            if rvo_math.abs_sq(self.simulator_.agents_[i].position_ - self.goals_[i]
                               ) > self.simulator_.agents_[i].dist_update_now_goal**2:
                return False

        return True


def main():
    if RVO_RENDER:
        pygame.init()
        screen = pygame.display.set_mode((750, 750), pygame.RESIZABLE)
        clock = pygame.time.Clock()

    blocks = Blocks()

    # Set up the scenario.
    blocks.setup_scenario()

    # Perform (and manipulate) the simulation.
    running = True
    bounds = (-10, 810, -310, 510)                 # x_min, x_max, y_min, y_max
    step = 0
    while running and not blocks.reached_goal() and not rospy.is_shutdown():
        if RVO_RENDER:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                elif event.type == pygame.VIDEORESIZE:
                    screen = pygame.display.set_mode((event.w, event.h), pygame.RESIZABLE)

            width, height = screen.get_size()
            blocks.update_visualization(screen, width, height, bounds)

            pygame.display.flip()
            clock.tick(60)  # 控制帧率为60帧每秒

        blocks.set_preferred_velocities()
        blocks.simulator_.step()

        step += 1
        print("current step: ", step)

    pygame.quit()


if __name__ == '__main__':
    main()
